from collections.abc import Sequence
from typing import TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, SkipExhausted

AnySample = TypeVar('AnySample')


class SequenceDatasetIterator(DatasetIterator[AnySample]):
    def __init__(self, sequence: Sequence[AnySample], state: int | None = None):
        self._sequence = sequence
        self._index = state if state is not None else 0

    def __next__(self) -> AnySample:
        if self._index >= len(self._sequence):
            raise StopIteration
        sample = self._sequence[self._index]
        self._index += 1
        return sample

    def __skip__(self, samples_n: int):
        remaining_samples_n = len(self._sequence) - self._index
        if samples_n <= remaining_samples_n:
            self._index += samples_n
        else:
            skipped_samples_n = remaining_samples_n
            self._index += skipped_samples_n
            raise SkipExhausted(skipped_samples_n)

    def __pause__(self) -> int:
        return self._index


class SequenceDataset(Dataset[AnySample]):
    def __init__(self, sequence: Sequence[AnySample]):
        self._sequence = sequence

    @property
    def base(self) -> Sequence[AnySample]:
        return self._sequence

    def __len__(self) -> int:
        return len(self._sequence)

    def __getitem__(self, index: int, /) -> AnySample:
        return self._sequence[index]

    def __iter__(self) -> DatasetIterator[AnySample]:
        return SequenceDatasetIterator(self._sequence)

    def __resume__(self, state: int | None = None, /) -> DatasetIterator[AnySample]:
        return SequenceDatasetIterator(self._sequence, state)
